import streamlit as st
import pandas as pd
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from PIL import Image
import json
import matplotlib.pyplot as plt
import os
from datetime import datetime

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 确保相册目录存在
if not os.path.exists("album"):
    os.makedirs("album")
for category in ["可回收物", "有害垃圾", "厨余垃圾", "其他垃圾"]:
    if not os.path.exists(f"album/{category}"):
        os.makedirs(f"album/{category}")

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# ========== 1. 加载模型 ==========
classifier = pipeline(Tasks.image_classification,
                      model='iic/cv_convnext-base_image-classification_garbage')

# ========== 2. 加载知识库 ==========
with open("knowledge.json", "r", encoding="utf-8") as f:
    knowledge = json.load(f)

# 初始化统计和历史记录
if "stats" not in st.session_state:
    st.session_state["stats"] = {"可回收物": 0, "有害垃圾": 0, "厨余垃圾": 0, "其他垃圾": 0}
if "history" not in st.session_state:
    st.session_state["history"] = []  # 存储历史记录：(文件名, 类别, 时间)

# 分类对应的投放建议
guide = {
    "可回收物": "请投放到 ♻️ 可回收物垃圾桶。",
    "有害垃圾": "⚠️ 请投放到 🔴 有害垃圾桶。",
    "厨余垃圾": "🍂 请投放到 🟢 厨余垃圾桶。",
    "其他垃圾": "🗑️ 请投放到 ⚪ 其他垃圾桶。"
}


# 保存图片到对应类别文件夹
def save_image(image, category):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}.jpg"
    filepath = f"album/{category}/{filename}"
    image.save(filepath)
    # 记录历史
    st.session_state["history"].append({
        "filename": filename,
        "category": category,
        "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "path": filepath
    })
    return filename


# 获取分类相册中的图片
def get_album_images(category):
    image_dir = f"album/{category}"
    if not os.path.exists(image_dir):
        return []
    images = []
    for filename in os.listdir(image_dir):
        if filename.endswith(('.jpg', '.png', '.jpeg')):
            images.append(os.path.join(image_dir, filename))
    return sorted(images, key=os.path.getmtime, reverse=True)  # 按时间排序


# ========== 3. Streamlit 界面 ==========
st.set_page_config(page_title="垃圾分类智能助手", layout="wide")
st.title("♻️ 垃圾分类智能助手")

# 添加相册标签页
tab1, tab2, tab3, tab4 = st.tabs(["垃圾识别", "知识问答", "统计分析", "分类相册"])

# --- Tab1: 图片识别 ---
with tab1:
    st.header("上传垃圾图片")
    uploaded_file = st.file_uploader("请上传垃圾图片", type=["jpg", "png", "jpeg"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="上传的图片", use_container_width=True)

        # 模型预测
        result = classifier(image)
        label = result["labels"][0]
        # 分割出大类型
        big_label = label.split("-")[0]
        score = result["scores"][0]

        st.subheader(f"预测结果：{label}（置信度 {score:.2f}）")
        st.success(guide.get(big_label, "请合理投放到对应垃圾桶。"))

        # 保存图片到相册
        save_image(image, big_label)

        # 更新统计
        if label in st.session_state["stats"]:
            st.session_state["stats"][label] += 1

# --- Tab2: 分类知识 ---
with tab2:
    st.header("垃圾分类知识库")
    category = st.selectbox("选择垃圾类别", list(knowledge.keys()))
    st.write(f"📘 **{category}** 的投放要求：")
    for rule in knowledge[category]:
        st.markdown(f"- {rule}")

# --- Tab3: 统计分析 ---
with tab3:
    st.header("上传垃圾统计结果")
    stats = st.session_state["stats"]
    df = pd.DataFrame(list(stats.items()), columns=["类别", "数量"])

    fig, ax = plt.subplots()
    ax.bar(df["类别"], df["数量"])
    ax.set_ylabel("数量")
    ax.set_title("垃圾类别分布统计")
    st.pyplot(fig)

# --- Tab4: 分类相册 ---
with tab4:
    st.header("分类相册")
    selected_category = st.selectbox("选择查看的垃圾类别", list(knowledge.keys()), key="album_category")

    # 显示该类别下的所有图片
    images = get_album_images(selected_category)
    st.write(f"共 {len(images)} 张图片")

    # 网格布局展示图片
    cols = st.columns(3)  # 3列布局
    for i, img_path in enumerate(images):
        with cols[i % 3]:
            try:
                img = Image.open(img_path)
                st.image(img, caption=os.path.basename(img_path), use_container_width=True)
                # 删除图片按钮
                if st.button(f"删除 {os.path.basename(img_path)}", key=f"del_{i}"):
                    os.remove(img_path)
                    # 更新历史记录
                    st.session_state["history"] = [h for h in st.session_state["history"] if h["path"] != img_path]
                    # 刷新页面
                    st.rerun()
            except Exception as e:
                st.error(f"无法加载图片: {e}")